import base64
import os
proxy = 'http://localhost:8234'
os.environ['HTTP_PROXY'] = proxy
os.environ['HTTPS_PROXY'] = proxy

from openai import OpenAI


def get_api_key(file_path):
    with open(file_path, 'r') as file:
        return file.read().strip()
    
from PIL import Image

# -*- coding: utf-8 -*-

import json
import re

def markdown_to_json(markdown_str):
    # 移除Markdown语法中可能存在的标记，如代码块标记等
    if markdown_str.startswith("```json"):
        markdown_str = markdown_str[7:].strip("` \n\t").strip()
    elif markdown_str.startswith("```"):
        markdown_str = markdown_str[3:].strip("` \n\t").strip()
    # 如果在markdown_str开头的50个字符内能够发现```json
    elif markdown_str[:50].find("```json") != -1:
        start_index = markdown_str[:50].find("```json")
        markdown_str = markdown_str[start_index + 7:].strip("` \n\t").strip()
    elif markdown_str[:50].find("```") != -1:
        start_index = markdown_str[:50].find("```")
        markdown_str = markdown_str[start_index:].strip("` \n\t").strip()

    print("尝试解析")
    print(markdown_str[:12],"...",markdown_str[-12:])  # for debug purposes
    print()

    # 将字符串转换为JSON字典
    json_dict = json.loads(markdown_str)

    return json_dict
    
def forced_extract(input_str, keywords):
    result = {key: "" for key in keywords}

    for key in keywords:
        # 使用正则表达式来查找关键词-值对
        pattern = f'"{key}":\s*"(.*?)"'
        match = re.search(pattern, input_str)
        if match:
            result[key] = match.group(1)
    
    return result


def forced_extract_guide(input_str):
    result = ""
    pattern = r'"guide":\s*"(.*?)"\s*}'
    match = re.search(pattern, input_str, re.DOTALL)
    if match:
        result = match.group(1)
    return result

def post_extract(input_str):
    keywords = ["analysis","guide"]

    result = None

    try:
        # 尝试使用 markdown_to_json 提取
        result = markdown_to_json(input_str)
        
        # 检查是否包含所有关键词
        if "guide" in result:
            return result
        
    except json.JSONDecodeError:
        # 如果 markdown_to_json 提取失败，则尝试使用 forced_extract 提取
        result = forced_extract(input_str, keywords)
        # result["guide"] = forced_extract_guide(input_str)

    return result




# Function to encode the image to base64
def encode_image(image_path):

    deal_path = image_path

    # 如果image对应的高度超过768，把其缩小到768, 并保存到lulu_exp/VLM/temp.jpg
    img = Image.open(image_path)

    # if image_path is a RGBA image, convert it to RGB
    if img.mode == 'RGBA':
        img = img.convert('RGB')
        img.save("lulu_exp/VLM/temp.jpg")
        img = Image.open("lulu_exp/VLM/temp.jpg")
        deal_path = "lulu_exp/VLM/temp.jpg"

    if img.height > 768:
        img = img.resize((int(img.width * 768 / img.height), 768))
        img.save("lulu_exp/VLM/temp.jpg")
        img = Image.open("lulu_exp/VLM/temp.jpg")
        deal_path = "lulu_exp/VLM/temp.jpg"


    with open(deal_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

# Function to query with image
def query_with_image(local_file_path, prompt, suffix_prompt = ""):
    # Initialize the OpenAI client
    client = OpenAI(api_key=get_api_key("lulu_exp/VLM/qwen_VLM_key.txt"), base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
    
    # Encode the image to base64
    base64_image = encode_image(local_file_path)

    messages = [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
            ]
    
    if suffix_prompt != "":
        messages.append({"type": "text", "text": suffix_prompt})

    # Send the request to the API
    response = client.chat.completions.create(
        model="qwen-vl-max-latest",
        messages=[{
            "role": "user",
            "content": messages
        }],
        response_format={
            'type': 'json_object'
        }
    )
    
    # Return the response
    return response.choices[0].message.content

# Test the function when the script is executed directly
if __name__ == "__main__":
    # image_path = "lulu_exp/VLM/dog_and_girl.jpeg"
    # prompt = "这是什么"
    
    # response = query_with_image(image_path, prompt)
    # print(response)

    response = """```json
{
  "analysis": "小朋友将数字3放到了数字4的位置。",
  "guide": "数字3不对哦，想想看，2后面应该是几呢？"
}
```"""

    print(post_extract(response))